from celery import Celery
from models.image_task import ImageTask
from models.ai_config import AIConfig
from app import create_app, db
import requests
import json
import base64
import os
from PIL import Image
import io

# 创建Celery实例
celery = Celery('meowai')

# 配置Celery
celery.conf.update(
    broker_url='redis://localhost:6379/0',
    result_backend='redis://localhost:6379/0',
    task_serializer='json',
    accept_content=['json'],
    result_serializer='json',
    timezone='Asia/Shanghai',
    enable_utc=True,
)

@celery.task
def generate_image_task(task_id):
    """异步图片生成任务"""
    app = create_app()
    with app.app_context():
        try:
            # 获取任务
            task = ImageTask.query.get(task_id)
            if not task:
                return {'error': '任务不存在'}
            
            # 更新任务状态为处理中
            task.update_status('processing', 10)
            
            # 获取AI配置
            ai_config = AIConfig.query.get(task.ai_config_id)
            if not ai_config or not ai_config.is_active:
                task.update_status('failed', error_message='AI服务不可用')
                return {'error': 'AI服务不可用'}
            
            # 根据不同的AI服务提供商调用相应的API
            if ai_config.provider == 'stable_diffusion':
                result = call_stable_diffusion_api(task, ai_config)
            elif ai_config.provider == 'baidu':
                result = call_baidu_api(task, ai_config)
            elif ai_config.provider == 'aliyun':
                result = call_aliyun_api(task, ai_config)
            else:
                # 模拟生成（用于演示）
                result = simulate_image_generation(task, ai_config)
            
            if result['success']:
                task.generated_image_url = result['image_url']
                task.update_status('completed', 100)
                return {'success': True, 'image_url': result['image_url']}
            else:
                task.update_status('failed', error_message=result['error'])
                return {'error': result['error']}
                
        except Exception as e:
            # 更新任务状态为失败
            task.update_status('failed', error_message=str(e))
            return {'error': str(e)}

def call_stable_diffusion_api(task, ai_config):
    """调用Stable Diffusion API"""
    try:
        # 更新进度
        task.update_status('processing', 30)
        
        # 读取原始图片并转换为base64
        image_path = os.path.join('uploads', os.path.basename(task.original_image_url))
        with open(image_path, 'rb') as f:
            image_base64 = base64.b64encode(f.read()).decode('utf-8')
        
        # 构建请求数据
        payload = {
            "text_prompts": [
                {
                    "text": task.positive_prompt,
                    "weight": 1.0
                },
                {
                    "text": task.negative_prompt,
                    "weight": -1.0
                }
            ],
            "init_image": image_base64,
            "init_image_mode": "IMAGE_STRENGTH",
            "image_strength": 0.35,
            "cfg_scale": task.cfg_scale,
            "samples": 1,
            "steps": task.steps,
            "width": task.width,
            "height": task.height,
        }
        
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {ai_config.api_key}"
        }
        
        # 更新进度
        task.update_status('processing', 60)
        
        # 发送请求
        response = requests.post(ai_config.api_endpoint, json=payload, headers=headers)
        
        if response.status_code == 200:
            response_data = response.json()
            
            # 解析生成的图片
            if 'artifacts' in response_data and len(response_data['artifacts']) > 0:
                image_data = response_data['artifacts'][0]['base64']
                
                # 保存生成的图片
                image_filename = f"generated_{task.id}.png"
                image_path = os.path.join('uploads', image_filename)
                
                with open(image_path, 'wb') as f:
                    f.write(base64.b64decode(image_data))
                
                return {
                    'success': True,
                    'image_url': f'/uploads/{image_filename}'
                }
            else:
                return {'success': False, 'error': '生成失败，返回数据异常'}
        else:
            return {'success': False, 'error': f'API调用失败: {response.text}'}
            
    except Exception as e:
        return {'success': False, 'error': str(e)}

def call_baidu_api(task, ai_config):
    """调用百度文心一格API"""
    try:
        # TODO: 实现百度API调用逻辑
        # 这里暂时返回模拟结果
        return simulate_image_generation(task, ai_config)
    except Exception as e:
        return {'success': False, 'error': str(e)}

def call_aliyun_api(task, ai_config):
    """调用阿里通义万相API"""
    try:
        # TODO: 实现阿里云API调用逻辑
        # 这里暂时返回模拟结果
        return simulate_image_generation(task, ai_config)
    except Exception as e:
        return {'success': False, 'error': str(e)}

def simulate_image_generation(task, ai_config):
    """模拟图片生成（用于演示）"""
    try:
        import time
        
        # 模拟处理时间
        for progress in [40, 60, 80]:
            task.update_status('processing', progress)
            time.sleep(1)  # 模拟处理时间
        
        # 复制原图作为"生成"结果（仅用于演示）
        original_path = os.path.join('uploads', os.path.basename(task.original_image_url))
        generated_filename = f"generated_{task.id}.jpg"
        generated_path = os.path.join('uploads', generated_filename)
        
        # 简单地复制文件（实际应该是AI生成的新图片）
        with open(original_path, 'rb') as src, open(generated_path, 'wb') as dst:
            dst.write(src.read())
        
        return {
            'success': True,
            'image_url': f'/uploads/{generated_filename}'
        }
        
    except Exception as e:
        return {'success': False, 'error': str(e)}

# 启动Celery Worker的命令
# celery -A tasks.image_generation.celery worker --loglevel=info